#include <vector>
#include <stack>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <algorithm>
#include <gflags/gflags.h>

#include "core/grape.h"
#include "utils/util.h"
#include "SimTemporal.hpp"

using namespace std;

template<class VD, class ED, class MessageT, class ResultT, class TemporalT>
class SimWorker: public Worker<VD, ED, MessageT, ResultT, TemporalT> {
	void pEval(Fragment<VD, ED> &fragment, MessageBuffer<MessageT> &messageBuffer, ResultT &partialResult, TemporalT &temporal) {
		unordered_map<VertexID, unordered_set<VertexID>> simSet, rmSet;

		unordered_map<VertexID, unordered_set<VertexID>> tSet;
		unordered_map<VertexID, unordered_set<VertexID>> prePattern;

		unordered_map<int, VertexID> tag2ID;
		unordered_map<VertexID, vector<unsigned>> cnt;
		unordered_map<VertexID, bool> active;

		unordered_set<VertexID> preV;

		auto allVertices = *fragment.getInnerVertices();

		for (auto v : *fragment.getOuterVertices()) {
			allVertices.insert(v);
		}

//		this->pattern.buildSignature(this->pattern.getAllVerticesID(), true, false);
		this->query.buildSignature();



		for (VertexID u = 0; u < this->query.getNumVertices(); ++u) {
			int tag = this->query.getVertexAttr(u);
			tag2ID[tag] = u;

			auto pSet = this->query.getParentsID(u);
			prePattern[u] = unordered_set<VertexID>(pSet.begin(), pSet.end());
			simSet[u] = unordered_set<VertexID>();
			rmSet[u] = unordered_set<VertexID>();
		}


		for (auto v : allVertices) {
			int tag = fragment.getVertexAttr(v);
			if (tag2ID.find(tag) != tag2ID.end()) {
//				if (fragment.isOuterVertex(v) || match(this->pattern.getSig(OUT_R1, tag2ID[tag]), fragment.getSig(OUT_R1, v))) {
//				if (!fragment.isInnerVertex(v) || match(this->query.getSignature(OUT_R1, tag2ID[tag]), fragment.getSignature(OUT_R1, v))) {
				cnt[v] = vector<unsigned>(this->query.getNumVertices(), 0);
				active[v] = true;
				if (fragment.hasChild(v)) {
					preV.insert(v);
				}
			}
		}

		for (auto vtx : active) {
			// vtx is a element of the active map.
			// vtx.first should be one of the active vertexID.
			VertexID v = vtx.first;
			int tag = fragment.getVertexAttr(v);
			VertexID u = tag2ID[tag];
			if (!this->query.hasChild(u) || fragment.hasChild(v)  || !fragment.isInnerVertex(v)) {
				for (auto p : fragment.getParentsID(v)) {
					if (active.find(p) != active.end()) {
						++cnt[p][tag];
						rmSet[u].insert(p);
					}
				}
			} else {
				active[v] = false;
			}
		}

		for (VertexID u = 0; u < this->query.getNumVertices(); ++u) {
			rmSet[u] = diff(preV, rmSet[u]);
		}

		queue<VertexID> que;
		vector<bool> inQue(this->query.getNumVertices(), false);
		for (VertexID u = 0; u < this->query.getNumVertices(); ++u) {
			if (!rmSet[u].empty()) {
				que.push(u);
				inQue[u] = true;
			}
		}

		for (; !que.empty(); que.pop()) {
			VertexID u = que.front();
			inQue[u] = false;
			for (auto w : rmSet[u]) {
				int tag = fragment.getVertexAttr(w);
				if (tag2ID.find(tag) != tag2ID.end()) {
					VertexID u1 = tag2ID[tag];
					if (prePattern[u].find(u1) != prePattern[u].end() && active[w]) {
						active[w] = false;
						for (auto w1 : fragment.getParentsID(w)) {
							if (active[w1]) {
								--cnt[w1][tag];
								if (cnt[w1][tag] == 0) {
									rmSet[u1].insert(w1);
									if (!inQue[u1]) {
										que.push(u1);
										inQue[u1] = true;
									}
								}
							}
						}
					}
				}
			}
			rmSet[u].clear();
		}

		for (auto vtx : active) {
			if (vtx.second) {
				VertexID v = vtx.first;
				int tag = fragment.getVertexAttr(v);
				VertexID u = tag2ID[tag];
				simSet[u].insert(v);
			}
		}

		partialResult = simSet;
		temporal = SimTemporal(tSet, prePattern, tag2ID, cnt, active);

		for (auto v : *fragment.getInnerVertices()) {
			if (!fragment.isBorderVertexWithIncomingEdge(v)) {
				continue;
			}
			int tag = fragment.getVertexAttr(v);
			if (!active[v] && tag2ID.find(tag) != tag2ID.end()) {
				auto dest = fragment.getMsgThroughIncomingDest(v);
				for (int fid = 0; fid < dest.size(); ++fid) {
					if (dest.test(fid)) {
						messageBuffer.add_message(fid, v);
					}
				}
			}
		}
	}

	void incEval(Fragment<VD, ED> &fragment, MessageBuffer<MessageT> &messageBuffer, ResultT &partialResult, TemporalT &temporal) {
		auto xIn = messageBuffer.get_messages();
		// auto Set = temporal.preSet;
		auto prePattern = temporal.prePattern;
		auto tag2ID = temporal.tag2ID;
		auto cnt = temporal.cnt;
		auto active = temporal.active;

		queue<VertexID> simQue;
		for (auto v : xIn) {
			if (active[v]) {
				simQue.push(v);
				active[v] = false;
			}
		}
		unordered_map<VertexID, unordered_set<VertexID>> deltaSet;
		for (int u = 0; u < this->query.getNumVertices(); ++u) {
			deltaSet[u] = unordered_set<VertexID>();
		}

		for (; !simQue.empty(); simQue.pop()) {
			auto v = simQue.front();

			if (fragment.isBorderVertexWithIncomingEdge(v)) {
				auto dest = fragment.getMsgThroughIncomingDest(v);
				for (int fid = 0; fid < dest.size(); ++fid) {
					if (dest.test(fid)) {
						messageBuffer.add_message(fid, v);
					}
				}
			}

			auto tag = fragment.getVertexAttr(v);
			auto u = tag2ID[tag];
			deltaSet[u].insert(v);
			for (auto v1 : fragment.getParentsID(v)) {
				if (active[v1]) {
					--cnt[v1][tag];
					if (cnt[v1][tag] == 0) {
						auto u1 = tag2ID[fragment.getVertexAttr(v1)];
						if (prePattern[u].find(u1) != prePattern[u].end()) {
							active[v1] = false;
							simQue.push(v1);
						}
					}
				}
			}
		}

		for (int u = 0; u < this->query.getNumVertices(); ++u) {
			partialResult[u] = diff(partialResult[u], deltaSet[u]);
		}

		temporal.cnt = cnt;
		temporal.active = active;
	}

	void assemble(vector<ResultT> &partialResults) {
		unordered_map<VertexID, unordered_set<VertexID>> finalResult;
		for (auto resResult : partialResults) {
			for (auto simU : resResult) {
				auto u = simU.first;
				for (auto v : simU.second) {
					finalResult[u].insert(v);
				}
			}
		}
		int tot = 0;

		if (this->result_file != "") {
			std::ofstream fout;
			fout.open(this->result_file);
			for (auto p : finalResult) {
				auto u = p.first;
				for (auto v : p.second) {
					fout << v << " -> " << u << endl;
				}
				tot += p.second.size();
			}
			fout.close();
		}

		LOG(INFO) << "Total matches: " << tot << std::endl;
	}

	template<class T>
	unordered_set<T> diff(const unordered_set<T> &a, const unordered_set<T>& b) {
		unordered_set<T> ret;
		for (auto ele : a) {
			if (b.find(ele) == b.end()) {
				ret.insert(ele);
			}
		}
		return ret;
	}
};

int main(int argc, char *argv[]) {

	FLAGS_stderrthreshold = 0;

	google::SetUsageMessage("Usage: mpiexec.hydra [mpi_opts] sim [grape_opts]");
	google::ParseCommandLineFlags(&argc, &argv, true);
	init_workers();

	WorkerParams param;

	param.algorithm = "simulation";
	param.use_pattern = true;
	param.enable_assemble = true;
	param.graph_load_strategy = ONLY_IN;

	google::ShutDownCommandLineFlags();

	SimWorker<int, int, VertexID, unordered_map<VertexID, unordered_set<VertexID>>, SimTemporal> simWorker;
	google::InitGoogleLogging("sim");
	simWorker.run(param);
	google::ShutdownGoogleLogging();

	worker_finalize();
	return 0;
}
